# -*- coding: utf-8 -*-
import sys
import os
import warnings
import numpy as np
import pandas as pd
from scipy.signal import welch, find_peaks
from tqdm import tqdm

# 如果你有自定义 pymagic 路径，加入 sys.path
dirs = ['/database/home/duansizhang/pywave']
for d in dirs:
    if d not in sys.path:
        sys.path.insert(0, d)

from pymagic.tools.ecgdllinterface import DataSource, EcgAnalyzer
from hrvanalysis import (
    remove_outliers, interpolate_nan_values, remove_ectopic_beats,
    get_time_domain_features, get_frequency_domain_features,
    get_geometrical_features
)

warnings.filterwarnings("ignore")

# ========== 参数配置 ==========
raw_file = '/database/private/mgcdb/raw/0a0b34b8220644848209d7a199afd6fc_1699056000000.raw'
out_dir  = '/database/home/duansizhang/hrv_predict/result'
fs_ecg   = 250       # 心电采样率 (Hz)
fs_acc   = 10        # 加速度采样率 (Hz)
win_sec  = 300       # 窗口长度：5分钟
nperseg_acc = fs_acc * 10  # 10秒段用于 Welch 频谱估计

os.makedirs(out_dir, exist_ok=True)

# --- 加速度数据读取 ---
hdata = DataSource(raw_file, 3)
acc = hdata.GetAccData()  # Nx3 array
N_acc = acc.shape[0]
window_len_acc = win_sec * fs_acc
num_windows = int(np.floor(N_acc / window_len_acc))

# --- R 波检测与 RR 预处理 ---
hdata_ecg = DataSource(raw_file, 3)
res = EcgAnalyzer(hdata_ecg)
rpos_all = np.array(res['rpos'], dtype=int)
anntype = np.array(res['anntype'], dtype=int)
rpos = rpos_all[anntype == 1]
if len(rpos) < 2:
    raise RuntimeError("R 波太少，无法计算 RR")
times = rpos / fs_ecg
rr_ms = np.diff(rpos) * 1000.0 / fs_ecg
rr_times = times[1:]
rr_clean = remove_outliers(rr_intervals=rr_ms, low_rri=300, high_rri=2000)
rr_interp = interpolate_nan_values(rr_intervals=rr_clean, interpolation_method='linear')
nn_beats = remove_ectopic_beats(rr_intervals=rr_interp, method='malik')
nn_interp = interpolate_nan_values(rr_intervals=nn_beats)
mask = ~np.isnan(nn_interp)
nn_array = np.array(nn_interp)[mask]
nn_times = rr_times[mask]
if len(nn_array) < 3:
    raise RuntimeError("有效 NN 段太短")

# --- 特征提取函数 ---
def extract_hrv_features(seg):
    feats = {}
    if len(seg) < 3:
        return None
    td = get_time_domain_features(seg)
    feats.update({k: float(td[k]) if td[k] is not None else np.nan for k in td})
    fd = get_frequency_domain_features(seg)
    feats.update({k: float(fd[k]) if fd[k] is not None else np.nan for k in fd})
    geo = get_geometrical_features(seg)
    feats.update({k: float(geo[k]) if geo[k] is not None else np.nan for k in geo})
    feats['nbeats'] = len(seg)
    return feats

# 结果行列表
overall_rows = []

# 对于每个窗口，提取 HRV 和加速度特征
for wi in tqdm(range(num_windows), desc="Processing windows"):
    # 时间窗口定义
    ws = wi * win_sec
    we = ws + win_sec

    # --- HRV 特征 ---
    seg_idx = (nn_times >= ws) & (nn_times < we)
    hr_seg = nn_array[seg_idx]
    hrv_feats = extract_hrv_features(hr_seg)
    if hrv_feats is None:
        continue

    # --- 加速度特征 ---
    start_acc = wi * window_len_acc
    end_acc = start_acc + window_len_acc
    seg_acc = acc[start_acc:end_acc, :]

    sma = np.sum(np.abs(seg_acc)) / window_len_acc
    std_mean = np.std(seg_acc, axis=0).mean()
    peak_total = sum(len(find_peaks(seg_acc[:, i])[0]) for i in range(3))
    f, Pxx = welch(seg_acc, fs=fs_acc, axis=0, nperseg=nperseg_acc)
    fft_power = np.sum(Pxx, axis=0).mean()
    spec_centroid = (f[:, None] * Pxx).sum(axis=0).mean() / fft_power
    spec_flux = np.sqrt(np.sum(np.diff(Pxx, axis=0)**2, axis=0)).mean()

    acc_feats = {
        'sma': sma,
        'std_mean': std_mean,
        'peak_total': peak_total,
        'fft_power': fft_power,
        'spec_centroid': spec_centroid,
        'spec_flux': spec_flux
    }

    # 合并到一行
    row = {
        'window_start': ws,
        'window_end': we,
        'timestamp': ws + win_sec/2,
        **hrv_feats,
        **acc_feats
    }
    overall_rows.append(row)

# 保存结果
df_out = pd.DataFrame(overall_rows)
out_path = os.path.join(out_dir, os.path.basename(raw_file).replace('.raw', '_combined_features.csv'))
df_out.to_csv(out_path, index=False, encoding='utf-8-sig')
print(f"已生成合并特征文件：{out_path}")
